#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
小样例数据测试 - 验证数据引擎的小数据集优化功能
"""

import asyncio
import pandas as pd
import json
import os
from datetime import datetime
from data_engine.core.engine import DataEngine
from data_engine.utils.llm_adapter import LLMClient
from data_engine.utils.logger import get_logger, setup_logging

# 设置日志级别为DEBUG
setup_logging(level="DEBUG")

# 获取测试日志器
test_logger = get_logger("small_data_test")


async def test_single_row_data():
    """测试单行数据优化"""
    test_logger.info("=== 测试单行数据优化 ===")
    
    # 创建单行数据
    single_data = pd.DataFrame({
        'name': ['张三'],
        'age': [28],
        'department': ['技术部'],
        'salary': [12000],
        'performance': [85.5]
    })
    
    # 初始化引擎
    llm_client = LLMClient(
        model="openai/qwen-turbo",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
        api_key="sk-kkkkkkkkkkkkkkkkkkk"
    )
    
    # 启用调试模式
    engine = DataEngine(llm_client, small_data_threshold=50, enable_grouping_filtering=True, debug_mode=True)
    
    # 加载数据集
    engine.load_dataset(
        df=single_data,
        dataset_desc="员工信息",
        table_desc="单个员工的基本信息",
        field_desc={
            "name": "姓名",
            "age": "年龄",
            "department": "部门",
            "salary": "薪资（元）",
            "performance": "绩效分数"
        }
    )
    
    # 测试问题
    question = "显示员工信息"
    test_logger.info(f"问题: {question}")
    
    start_time = datetime.now()
    result = await engine.process_query(question)
    end_time = datetime.now()
    
    execution_time = (end_time - start_time).total_seconds()
    
    test_logger.info(f"✅ 处理时间: {execution_time:.3f}秒")
    test_logger.info(f"✅ 成功: {result.success}")
    test_logger.info(f"✅ 图表类型: {result.chart_type}")
    test_logger.info(f"✅ 消息: {result.message}")
    
    return {
        "test_type": "single_row",
        "success": result.success,
        "execution_time": execution_time,
        "chart_type": result.chart_type,
        "data_rows": len(single_data),
        "chart_config": result.chart_config
    }


async def test_small_dataset():
    """测试小数据集优化"""
    test_logger.info("\n=== 测试小数据集优化 ===")
    
    # 创建小数据集（5行）
    small_data = pd.DataFrame({
        'product': ['产品A', '产品B', '产品C', '产品D', '产品E'],
        'sales': [1200, 800, 1500, 900, 1100],
        'category': ['电子', '服装', '电子', '食品', '服装'],
        'price': [299, 89, 599, 25, 129]
    })
    
    # 初始化引擎（禁用分组过滤以触发小数据集优化）
    llm_client = LLMClient(
        model="openai/qwen-turbo",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
        api_key="sk-kkkkkkkkkkkkkkkkkkk"
    )
    
    # 启用调试模式
    engine = DataEngine(llm_client, small_data_threshold=50, enable_grouping_filtering=False, debug_mode=True)
    
    # 加载数据集
    engine.load_dataset(
        df=small_data,
        dataset_desc="产品销售数据",
        table_desc="各产品的销售情况",
        field_desc={
            "product": "产品名称",
            "sales": "销售数量",
            "category": "产品类别",
            "price": "单价（元）"
        }
    )
    
    # 测试问题
    question = "显示产品销售情况"
    test_logger.info(f"问题: {question}")
    
    start_time = datetime.now()
    result = await engine.process_query(question)
    end_time = datetime.now()
    
    execution_time = (end_time - start_time).total_seconds()
    
    test_logger.info(f"✅ 处理时间: {execution_time:.3f}秒")
    test_logger.info(f"✅ 成功: {result.success}")
    test_logger.info(f"✅ 图表类型: {result.chart_type}")
    test_logger.info(f"✅ 消息: {result.message}")
    
    return {
        "test_type": "small_dataset",
        "success": result.success,
        "execution_time": execution_time,
        "chart_type": result.chart_type,
        "data_rows": len(small_data),
        "chart_config": result.chart_config
    }


async def test_medium_dataset():
    """测试中等数据集（标准流程）"""
    test_logger.info("\n=== 测试中等数据集（标准流程） ===")
    
    # 创建中等数据集（100行）
    import numpy as np
    np.random.seed(42)
    
    medium_data = pd.DataFrame({
        'month': [f'2024-{i:02d}' for i in range(1, 13)] * 8 + [f'2024-{i:02d}' for i in range(1, 5)],
        'department': ['销售部', '技术部', '市场部', '财务部'] * 25,
        'sales_amount': np.random.randint(50000, 200000, 100),
        'region': ['北京', '上海', '广州', '深圳', '杭州'] * 20
    })
    
    # 初始化引擎
    llm_client = LLMClient(
        model="openai/qwen-turbo",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
        api_key="sk-kkkkkkkkkkkkkkkkkkk"
    )
    
    # 启用调试模式
    engine = DataEngine(llm_client, small_data_threshold=50, enable_grouping_filtering=True, debug_mode=True)
    
    # 加载数据集
    engine.load_dataset(
        df=medium_data,
        dataset_desc="公司销售数据",
        table_desc="各部门各地区每月销售金额",
        field_desc={
            "month": "月份",
            "department": "部门",
            "sales_amount": "销售金额（元）",
            "region": "地区"
        }
    )
    
    # 测试问题
    question = "显示各部门的销售趋势"
    test_logger.info(f"问题: {question}")
    
    start_time = datetime.now()
    result = await engine.process_query(question)
    end_time = datetime.now()
    
    execution_time = (end_time - start_time).total_seconds()
    
    test_logger.info(f"✅ 处理时间: {execution_time:.3f}秒")
    test_logger.info(f"✅ 成功: {result.success}")
    test_logger.info(f"✅ 图表类型: {result.chart_type}")
    test_logger.info(f"✅ 消息: {result.message}")
    
    return {
        "test_type": "medium_dataset",
        "success": result.success,
        "execution_time": execution_time,
        "chart_type": result.chart_type,
        "data_rows": len(medium_data),
        "chart_config": result.chart_config
    }


async def main():
    """主测试函数"""
    test_logger.info("🚀 开始小样例数据测试...\n")
    
    results = []
    
    try:
        # 测试单行数据
        result1 = await test_single_row_data()
        results.append(result1)
        
        # 测试小数据集
        result2 = await test_small_dataset()
        results.append(result2)
        
        # 测试中等数据集（对比）
        result3 = await test_medium_dataset()
        results.append(result3)
        
        # 保存测试结果
        test_summary = {
            "timestamp": datetime.now().isoformat(),
            "total_tests": len(results),
            "successful_tests": sum(1 for r in results if r["success"]),
            "results": results,
            "performance_comparison": {
                "single_row_time": results[0]["execution_time"],
                "small_dataset_time": results[1]["execution_time"],
                "medium_dataset_time": results[2]["execution_time"],
                "optimization_ratio": {
                    "single_vs_medium": results[2]["execution_time"] / results[0]["execution_time"] if results[0]["execution_time"] > 0 else 0,
                    "small_vs_medium": results[2]["execution_time"] / results[1]["execution_time"] if results[1]["execution_time"] > 0 else 0
                }
            }
        }
        
        # 确保输出目录存在
        os.makedirs("test_output", exist_ok=True)
        
        # 保存结果
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"test_output/small_data_test_{timestamp}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(test_summary, f, ensure_ascii=False, indent=2)
        
        test_logger.info(f"\n📊 测试总结:")
        test_logger.info(f"✅ 成功测试: {test_summary['successful_tests']}/{test_summary['total_tests']}")
        test_logger.info(f"⚡ 单行数据处理时间: {results[0]['execution_time']:.3f}秒")
        test_logger.info(f"⚡ 小数据集处理时间: {results[1]['execution_time']:.3f}秒")
        test_logger.info(f"⚡ 中等数据集处理时间: {results[2]['execution_time']:.3f}秒")
        test_logger.info(f"📈 性能提升倍数:")
        test_logger.info(f"   单行 vs 中等: {test_summary['performance_comparison']['optimization_ratio']['single_vs_medium']:.1f}x")
        test_logger.info(f"   小集 vs 中等: {test_summary['performance_comparison']['optimization_ratio']['small_vs_medium']:.1f}x")
        test_logger.info(f"📁 详细结果已保存到: {filename}")
        
        # 输出ECharts配置的JSON - 统一格式
        for i, result in enumerate(results):
            if result.get('chart_config'):
                # 创建统一的图表配置格式
                chart_output = {
                    "metadata": {
                        "test_type": result["test_type"],
                        "chart_type": result["chart_type"],
                        "data_rows": result["data_rows"],
                        "execution_time": result["execution_time"],
                        "timestamp": datetime.now().isoformat(),
                        "source_file": "small_data_test.py"
                    },
                    "chart_config": result['chart_config'],
                    "format_version": "1.0"
                }
                
                chart_filename = f"test_output/small_data_test_chart_{i+1}_{timestamp}.json"
                with open(chart_filename, 'w', encoding='utf-8') as f:
                    json.dump(chart_output, f, ensure_ascii=False, indent=2)
                test_logger.info(f"📊 图表配置已保存到: {chart_filename}")
        
    except Exception as e:
        test_logger.error(f"❌ 测试失败: {str(e)}")
        import traceback
        traceback.print_exc()
        
        # 保存错误信息
        error_data = {
            "timestamp": datetime.now().isoformat(),
            "error": str(e),
            "traceback": traceback.format_exc()
        }
        
        os.makedirs("test_output", exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"test_output/small_data_test_error_{timestamp}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(error_data, f, ensure_ascii=False, indent=2)
        
        test_logger.error(f"📁 错误信息已保存到: {filename}")


if __name__ == "__main__":
    asyncio.run(main())